#!/usr/bin/env python3
"""
Preprocess the ETH3D dataset (supporting PFM depth format).

This script processes images from the ETH3D dataset, reading camera parameters from
the `_cam.txt` files, loading and scaling depth maps, resizing images, and saving
the processed RGB images, depth maps, and camera metadata into an output directory
structured to match the Hypersim processed format.

It relies on the `cams/index2prefix.txt` file within each scene to establish the
correct correspondence between images and camera parameter files.

This script is designed to be used with an existing `utils` directory that provides
helper functions like `crop_resize_if_necessary`.

Usage:
    python preprocess_eth3d.py --eth3d_dir /path/to/eth3d \
                               --output_dir /path/to/processed_eth3d
"""
import sys
sys.path.append('D:/codes/working/3D/a3R')
sys.path.append('D:/codes/working/3D/a3R/src')
# sys.path.append('/lc/code/3D/a3R')
# sys.path.append('/lc/code/3D/a3R/src')
import os
import os.path as osp
import argparse
import re
from tqdm import tqdm
import numpy as np
import argparse
from PIL import Image

from src.dust3r.datasets.base.base_multiview_dataset import cropping
from utils.cropping import crop_resize_if_necessary
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
import cv2


def get_parser():
    """Defines the command-line arguments for the script."""
    parser = argparse.ArgumentParser(
        description="Preprocess the ETH3D dataset by converting camera parameters, "
        "resizing images, and saving processed outputs in a Hypersim-like format."
    )
    parser.add_argument(
        "--eth3d_dir",
        required=True,
        help="Root directory of the ETH3D dataset (containing scenes like 'forest')."
    )
    parser.add_argument(
        "--output_dir",
        default="data/eth3d_processed",
        help="Output directory for processed ETH3D data."
    )
    parser.add_argument(
        "--target_height",
        type=int,
        default=384,
        help="Target height to resize images and depth maps to."
    )
    parser.add_argument(
        "--target_width",
        type=int,
        default=687,
        help="Target width to resize images and depth maps to."
    )
    return parser

def main(db_root, output_dir, target_wh):
    """Main function to orchestrate the preprocessing."""
    print(">> Listing all scenes in the ETH3D dataset")
    scenes = []
    for f in sorted(os.listdir(db_root)):
        scene_path = osp.join(db_root, f)
        if osp.isdir(scene_path) and osp.isdir(osp.join(scene_path, 'cams')):
            scenes.append(f)
    
    assert scenes, f"Did not find any scenes at {db_root}"
    print(f"   (Found {len(scenes)} scenes)")

    for scene in tqdm(scenes, desc="Processing Scenes"):
        scene_root = osp.join(db_root, scene)
        cam_dir = osp.join(scene_root, 'cams')
        index_file = osp.join(cam_dir, 'index2prefix.txt')

        if not osp.exists(index_file):
            tqdm.write(f"!! Skipping scene '{scene}': index2prefix.txt not found.")
            continue

        with open(index_file) as f:
            lines = [line.strip() for line in f.readlines() if line.strip()]

        mappings = [line.split() for line in lines[1:]]

        for mapping in tqdm(mappings, desc=f"  Processing {scene}", leave=False):
            img_idx_str, relative_img_path = mapping
            cam_file_num = int(img_idx_str)

            # --- Define Subscene and Output Directory ---
            # Get subscene name from the image path, e.g., 'images_rig_cam4_undistorted' -> 'images_rig_cam4'
            subscene_name = osp.basename(osp.dirname(relative_img_path)).replace('_undistorted', '')
            subscene_outdir = osp.join(output_dir, scene, subscene_name)
            os.makedirs(subscene_outdir, exist_ok=True)

            # --- Define File Paths and Names ---
            cam_path = osp.join(cam_dir, f'{cam_file_num:08d}_cam.txt')
            img_path = osp.join(scene_root, 'images', relative_img_path)
            
            # Construct the .pfm depth path
            base_filename, _ = osp.splitext(osp.basename(relative_img_path))
            depth_rig_dir_name = subscene_name # e.g., 'images_rig_cam4'
            depth_path = osp.join(scene_root, 'depths', depth_rig_dir_name, base_filename + '.pfm')

            # Use the original filename (timestamp) as the base for output files
            output_name = base_filename
            
            process_single_view(
                img_path, depth_path, cam_path, subscene_outdir, output_name, target_wh
            )

    print(f">> Done, saved everything in {output_dir}/")

def process_single_view(img_path, depth_path, cam_path, out_dir, output_name, target_wh):
    """
    Processes a single view and saves all its files directly into the output directory.
    """
    # Define output paths directly in the target directory
    out_rgb_path = osp.join(out_dir, f"{output_name}_rgb.png")
    out_depth_path = osp.join(out_dir, f"{output_name}_depth.exr")
    out_cam_path = osp.join(out_dir, f"{output_name}_cam.npz")

    # Skip if final camera file already exists
    # if osp.isfile(out_cam_path):
    #     return

    if not all(osp.exists(p) for p in [img_path, depth_path, cam_path]):
        print(f'{depth_path} file not exists')
        return

    # Load data from disk
    intrinsics, T_cam2world = _load_pose(cam_path, ret_44=True)
    color_image = cv2.cvtColor(
        cv2.imread(img_path, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB
    )
    depth_map = load_pfm_file(depth_path)

    # Resize image, depth, and update intrinsics using the external utility
    image_out, depth_out, intrinsics_out = crop_resize_if_necessary(
        color_image, depth_map, intrinsics, target_wh
    )
    intrinsics, T_cam2world = cropping.get_center_camera(intrinsics, T_cam2world, depthmap=depth_out)
    # Save the processed data
    image_out.save(out_rgb_path)
    # cv2.imwrite(out_rgb_path, cv2.cvtColor(image_out, cv2.COLOR_RGB2BGR))
    cv2.imwrite(out_depth_path, depth_out.astype(np.float32))
    np.savez(
        out_cam_path,
        intrinsics=intrinsics_out.astype(np.float32),
        pose=T_cam2world.astype(np.float32),
    )

def _load_pose(path, ret_44=False):
    # f = open(path)
    with open(path, 'r') as f:
        RT = np.loadtxt(f, skiprows=1, max_rows=4, dtype=np.float32)
        assert RT.shape == (4, 4)
        RT = np.linalg.inv(RT)  # world2cam to cam2world

        K = np.loadtxt(f, skiprows=2, max_rows=3, dtype=np.float32)
        assert K.shape == (3, 3)

    if ret_44:
        return K, RT
    return K, RT[:3, :3], RT[:3, 3]

def load_pfm_file(file_path):
    with open(file_path, "rb") as file:
        header = file.readline().decode("UTF-8").strip()

        if header == "PF":
            is_color = True
        elif header == "Pf":
            is_color = False
        else:
            raise ValueError("The provided file is not a valid PFM file.")

        dimensions_line = file.readline().decode("UTF-8").strip()
        dimensions = re.match(r"^(\d+)\s+(\d+)$", dimensions_line)

        if dimensions:
            img_width, img_height = map(int, dimensions.groups())
        else:
            raise ValueError("Invalid PFM header format.")

        endian_scale = float(file.readline().decode("UTF-8").strip())
        if endian_scale < 0:
            dtype = "<f"  # little-endian
        else:
            dtype = ">f"  # big-endian

        try:
            data_buffer = file.read()
            img_data = np.frombuffer(data_buffer, dtype=dtype)
        except Exception as e:
            print(f"Error processing {file_path}: {e}")
            sys.exit(0)


        if is_color:
            img_data = np.reshape(img_data, (img_height, img_width, 3))
        else:
            img_data = np.reshape(img_data, (img_height, img_width))

        # PFM files are stored bottom-to-top, so we need to flip it.
        img_data = cv2.flip(img_data, 0)

    return img_data

if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()
    target_wh = (args.target_width, args.target_height)
    main(args.eth3d_dir, args.output_dir, target_wh)